#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/uio.h>
#include <errno.h>
#include <time.h>
#include <stdint.h>
#include <pthread.h>
#include <arpa/inet.h>
#include "hash.h"
#include "public.h"

// 原有哈希函数不变
static size_t hash_function(const char* key, size_t bucket_count) {
    size_t hash = 5381;
    int c;
    while ((c = *key++)) {
        hash = ((hash << 5) + hash) + c;  // hash * 33 + c
    }
    return hash % bucket_count;
}

// 1. 修改创建函数，添加锁初始化（与红黑树锁逻辑一致）
HashTable* hashtable_create(size_t bucket_count) {
    if (bucket_count == 0) return NULL;

    HashTable* ht = (HashTable*)kvs_malloc(sizeof(HashTable));
    
    if (!ht) return NULL;

    // 分配桶数组并初始化为NULL
    ht->buckets = (HashNode**)calloc(bucket_count, sizeof(HashNode*));
    if (!ht->buckets) {
        kvs_free(ht);
        return NULL;
    }

    ht->bucket_count = bucket_count;
    ht->count = 0;

    // 新增：初始化互斥锁（递归锁，支持同一线程重复加锁）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutexattr_t mutex_attr;
    pthread_mutexattr_init(&mutex_attr);
    pthread_mutexattr_settype(&mutex_attr, PTHREAD_MUTEX_RECURSIVE);
    pthread_mutex_init(&ht->mutex_lock, &mutex_attr);
    pthread_mutexattr_destroy(&mutex_attr);
#endif

    // 新增：初始化读写锁（写者优先）
#if HASH_TABLE_RW_LOCK
    pthread_rwlockattr_t rw_attr;
    pthread_rwlockattr_init(&rw_attr);
#ifdef PTHREAD_RWLOCK_PREFER_WRITER_NONRECURSIVE_NP
    pthread_rwlockattr_setkind_np(&rw_attr, PTHREAD_RWLOCK_PREFER_WRITER_NONRECURSIVE_NP);
#elif defined(PTHREAD_RWLOCK_PREFER_WRITER)
    pthread_rwlockattr_setkind(&rw_attr, PTHREAD_RWLOCK_PREFER_WRITER);
#endif
    pthread_rwlock_init(&ht->rw_lock, &rw_attr);
    pthread_rwlockattr_destroy(&rw_attr);
#endif

    return ht;
}

// 2. 新增：清空哈希表（保留结构，删除所有节点）
int hashtable_clear(HashTable* ht) {
    if (!ht) return -1;

    // 加锁（写操作）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_wrlock(&ht->rw_lock);
#endif

    // 遍历所有桶，释放节点
    for (size_t i = 0; i < ht->bucket_count; i++) {
        HashNode* current = ht->buckets[i];
        while (current) {
            HashNode* next = current->next;
            kvs_free(current->key);
            kvs_free(current->value);
            kvs_free(current);
            current = next;
        }
        ht->buckets[i] = NULL; // 重置桶头指针
    }
    ht->count = 0; // 重置元素计数

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return 0;
}

// 3. 原有操作函数添加锁（仅展示修改部分，逻辑不变）
int hashtable_insert(HashTable* ht, const char* key, const char* value) {
    if (!ht || !key || !value) return -1;

    // 加锁（写操作）
#if HASH_TABLE_MUTEX_LOCK

    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_wrlock(&ht->rw_lock);
#endif

    // 原有插入逻辑不变...
    size_t index = hash_function(key, ht->bucket_count);
    HashNode* current = ht->buckets[index];
    while (current) {
        if (strcmp(current->key, key) == 0) {
            // 解锁后返回
#if HASH_TABLE_MUTEX_LOCK
            pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
            pthread_rwlock_unlock(&ht->rw_lock);
#endif
            return 1;
        }
        current = current->next;
    }

    HashNode* new_node = (HashNode*)kvs_malloc(sizeof(HashNode));
    if (!new_node) {
#if HASH_TABLE_MUTEX_LOCK
        pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
        pthread_rwlock_unlock(&ht->rw_lock);
#endif
        return -1;
    }

    new_node->key = strdup(key);
    new_node->value = strdup(value);
    if (!new_node->key || !new_node->value) {
        kvs_free(new_node->key);
        kvs_free(new_node->value);
        kvs_free(new_node);
#if HASH_TABLE_MUTEX_LOCK
        pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
        pthread_rwlock_unlock(&ht->rw_lock);
#endif
        return -1;
    }

    new_node->next = ht->buckets[index];
    ht->buckets[index] = new_node;
    ht->count++;

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return 0;
}

int hashtable_delete(HashTable* ht, const char* key) {
    if (!ht || !key) return -1;

    // 加锁（写操作）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_wrlock(&ht->rw_lock);
#endif

    // 原有删除逻辑不变...
    size_t index = hash_function(key, ht->bucket_count);
    HashNode* current = ht->buckets[index];
    HashNode* prev = NULL;
    while (current) {
        if (strcmp(current->key, key) == 0) {
            if (prev) prev->next = current->next;
            else ht->buckets[index] = current->next;

            kvs_free(current->key);
            kvs_free(current->value);
            kvs_free(current);
            ht->count--;

            // 解锁后返回
#if HASH_TABLE_MUTEX_LOCK
            pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
            pthread_rwlock_unlock(&ht->rw_lock);
#endif
            return 0;
        }
        prev = current;
        current = current->next;
    }

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return 1;
}

char* hashtable_find(HashTable* ht, const char* key) {
    if (!ht || !key) return NULL;

    // 加锁（读操作）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_rdlock(&ht->rw_lock);
#endif

    // 原有查找逻辑不变...
    size_t index = hash_function(key, ht->bucket_count);
    HashNode* current = ht->buckets[index];
    char* result = NULL;
    while (current) {
        if (strcmp(current->key, key) == 0) {
            result = current->value;
            break;
        }
        current = current->next;
    }

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return result;
}

int hashtable_update(HashTable* ht, const char* key, const char* new_value) {
    if (!ht || !key || !new_value) return -1;

    // 加锁（写操作）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_wrlock(&ht->rw_lock);
#endif

    // 原有更新逻辑不变...
    size_t index = hash_function(key, ht->bucket_count);
    HashNode* current = ht->buckets[index];
    while (current) {
        if (strcmp(current->key, key) == 0) {
            char* updated_value = strdup(new_value);
            if (!updated_value) {
#if HASH_TABLE_MUTEX_LOCK
                pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
                pthread_rwlock_unlock(&ht->rw_lock);
#endif
                return -1;
            }
            kvs_free(current->value);
            current->value = updated_value;

#if HASH_TABLE_MUTEX_LOCK
            pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
            pthread_rwlock_unlock(&ht->rw_lock);
#endif
            return 0;
        }
        current = current->next;
    }

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return 1;
}

// 4. 修改销毁函数，添加锁销毁
void hashtable_destroy(HashTable* ht) {
    if (!ht) return;

    // 加锁（确保销毁时无并发操作）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_wrlock(&ht->rw_lock);
#endif

    // 原有销毁逻辑不变...
    for (size_t i = 0; i < ht->bucket_count; i++) {
        HashNode* current = ht->buckets[i];
        while (current) {
            HashNode* next = current->next;
            kvs_free(current->key);
            kvs_free(current->value);
            kvs_free(current);
            current = next;
        }
    }
    kvs_free(ht->buckets);

    // 销毁锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
    pthread_mutex_destroy(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
    pthread_rwlock_destroy(&ht->rw_lock);
#endif

    kvs_free(ht);
}

/**
 * 获取元素数量（线程安全版）
 * @param ht 哈希表
 * @return 元素数量，哈希表为NULL时返回0
 */
size_t hashtable_count(HashTable* ht) {
    if (!ht) return 0;

    // 加读锁（读取count变量）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_rdlock(&ht->rw_lock);
#endif

    size_t count = ht->count;  // 读取元素数量

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return count;
}



// ========================= 全量持久化 =============================
// 5. 新增：全量持久化相关函数（仿照红黑树）
// 辅助函数：遍历哈希表并写入文件
static int hashtable_traverse_persist(HashTable* ht, int fd) {
    if (!ht || fd == -1) return -1;

    // 遍历所有桶
    for (size_t i = 0; i < ht->bucket_count; i++) {
        HashNode* current = ht->buckets[i];
        while (current) {
            int key_len = strlen(current->key);
            int val_len = strlen(current->value);

#if 1 // MASTER_SLAVE 为了方便各种模式都能使用正常同一个文件，这里统一存储/读取方式
            // 网络字节序转换
            uint32_t net_key_len = htonl(key_len);
            uint32_t net_val_len = htonl(val_len);
            struct iovec iov[4] = {
                {&net_key_len, 4},
                {&net_val_len, 4},
                {current->key, key_len},
                {current->value, val_len}
            };
#else
            struct iovec iov[4] = {
                {&key_len, 4},
                {&val_len, 4},
                {current->key, key_len},
                {current->value, val_len}
            };
#endif

            ssize_t total = writev(fd, iov, 4);
            if (total != 4 + 4 + key_len + val_len) {
                perror("hashtable traverse persist failed");
                return -1;
            }
            current = current->next;
        }
    }
    return 0;
}

// 全量持久化
int hashtable_full_store(HashTable* ht, const char* full_filename) {
    if (!ht || !full_filename) return -1;

    // 加读锁（仅读取数据）
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_rdlock(&ht->rw_lock);
#endif

    int fd = open(full_filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (fd == -1) {
        perror("full store open failed");
#if HASH_TABLE_MUTEX_LOCK
        pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
        pthread_rwlock_unlock(&ht->rw_lock);
#endif
        return -1;
    }

    int ret = hashtable_traverse_persist(ht, fd);
    fsync(fd);
    close(fd);

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return ret;
}

// 带备份的全量持久化
int hashtable_full_store_with_backup(HashTable* ht, const char* full_filename) {
    if (!ht || !full_filename) return -1;

    // 加读锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_lock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_rdlock(&ht->rw_lock);
#endif

    // 备份旧文件
    struct stat st;
    if (stat(full_filename, &st) == 0) {
        char backup_filename[256];
        time_t now = time(NULL);
        struct tm* tm = localtime(&now);
        snprintf(backup_filename, sizeof(backup_filename),
                 "%s.%04d%02d%02d%02d%02d",
                 full_filename, tm->tm_year + 1900, tm->tm_mon + 1, tm->tm_mday,
                 tm->tm_hour, tm->tm_min);
        rename(full_filename, backup_filename);
    }

    // 写入新数据
    int fd = open(full_filename, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if (fd == -1) {
        perror("full store with backup open failed");
#if HASH_TABLE_MUTEX_LOCK
        pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
        pthread_rwlock_unlock(&ht->rw_lock);
#endif
        return -1;
    }

    int ret = hashtable_traverse_persist(ht, fd);
    fsync(fd);
    close(fd);

    // 解锁
#if HASH_TABLE_MUTEX_LOCK
    pthread_mutex_unlock(&ht->mutex_lock);
#elif HASH_TABLE_RW_LOCK
    pthread_rwlock_unlock(&ht->rw_lock);
#endif

    return ret;
}

// 从文件恢复哈希表
int hashtable_full_restore(const char* full_filename, HashTable* ht) {
    if (!full_filename || !ht) return -1;

    // 先清空现有数据
    if (hashtable_clear(ht) != 0) {
        fprintf(stderr, "clear before restore failed\n");
        return -1;
    }

    int fd = open(full_filename, O_RDONLY);
    if (fd == -1) {
        perror("restore open failed");
        return -1;
    }

#if 1 // MASTER_SLAVE 为了方便各种模式都能使用正常同一个文件，这里统一存储/读取方式
    uint32_t net_key_len, net_val_len;
#endif
    int key_len, val_len;
    char *key_buf = NULL, *val_buf = NULL;
    long int read_count = 0;
    int is_eof = 0;

    while (1) {
#if 1 // MASTER_SLAVE 为了方便各种模式都能使用正常同一个文件，这里统一存储/读取方式
        // 读网络字节序长度
        ssize_t n = read(fd, &net_key_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;
        key_len = ntohl(net_key_len);

        n = read(fd, &net_val_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;
        val_len = ntohl(net_val_len);
#else
        // 读主机字节序长度
        ssize_t n = read(fd, &key_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;

        n = read(fd, &val_len, 4);
        if (n == 0) { is_eof = 1; break; }
        if (n != 4) break;
#endif

        // 校验长度
        if (key_len <= 0 || val_len <= 0) {
            fprintf(stderr, "invalid key/val length\n");
            break;
        }

        // 动态分配缓冲区
        key_buf = kvs_malloc(key_len + 1);
        if (!key_buf) { perror("kvs_malloc key_buf failed"); break; }

        val_buf = kvs_malloc(val_len + 1);
        if (!val_buf) {
            perror("kvs_malloc val_buf failed");
            kvs_free(key_buf);
            key_buf = NULL;
            break;
        }

        // 读键值
        n = read(fd, key_buf, key_len);
        if (n == 0) { is_eof = 1; break; }
        if (n != key_len) break;
        key_buf[key_len] = '\0';

        n = read(fd, val_buf, val_len);
        if (n == 0) { is_eof = 1; break; }
        if (n != val_len) break;
        val_buf[val_len] = '\0';

        // 插入哈希表（内部已加锁）
        if (hashtable_insert(ht, key_buf, val_buf) == 0) {
            read_count++;
        }

        // 释放缓冲区
        kvs_free(key_buf);
        kvs_free(val_buf);
        key_buf = NULL;
        val_buf = NULL;
    }

    LOG("成功读取了%ld个节点\n", read_count);

    // 处理异常
    if (!is_eof && errno != 0 && errno != EINTR) {
        perror("restore read failed");
        hashtable_clear(ht); // 恢复失败则清空
    }

    // 清理资源
    if (key_buf) kvs_free(key_buf);
    if (val_buf) kvs_free(val_buf);
    close(fd);

    return 0;
}